package jtg.graphics;

import javafx.util.Pair;
import soot.*;
import soot.coffi.CFG;
import soot.jimple.InvokeExpr;
import soot.jimple.JimpleBody;
import soot.jimple.SwitchStmt;
import soot.jimple.internal.*;
import soot.options.Options;
import soot.toolkits.graph.*;
import soot.util.Chain;

import java.io.File;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

public class SootCFG {


    public static UnitGraph getMethodCFG(Body body){
        // 使用 ClassicCompleteUnitGraph 类构造一个完整的控制流图。
        // ClassicCompleteUnitGraph 是 Soot 提供的用于构建完整控制流图的类，其中包含了所有可能的边，即每个基本块都与其他所有基本块相连。
        UnitGraph ug = new ClassicCompleteUnitGraph(body);
        return ug;
    }
    public static UnitGraph getMethodCFG(String sourceDirectory,String clsName,String methodName){
        Body body = getMethodBody(sourceDirectory,clsName,methodName);
        // 使用 ClassicCompleteUnitGraph 类构造一个完整的控制流图。
        // ClassicCompleteUnitGraph 是 Soot 提供的用于构建完整控制流图的类，其中包含了所有可能的边，即每个基本块都与其他所有基本块相连。
        UnitGraph ug = new ClassicCompleteUnitGraph(body);
        return ug;
    }

    // 展开控制图，把调用的类的函数展开，并存到invokes里面
    public static void expandCFG(Body body, Set<Unit> invokes) {
        // 拿到body所有的units
        UnitPatchingChain units = body.getUnits();
        List<Unit> unitList = new ArrayList<>(units);
        // 遍历所有节点，有调用就扩展、
        for (int i = 0; i < unitList.size(); i++) {
            Unit unit = unitList.get(i);

            // 调用节点，不是赋值语句
            if (unit instanceof JInvokeStmt) {
                invokes.add(unit);
                // 类型转换一下，不然后面没法做
                JInvokeStmt jInvokeStmt = (JInvokeStmt) unit;
                // 拿到invoke的 类.方法
                InvokeExpr invokeExpr = jInvokeStmt.getInvokeExpr();
                // 去找调用的类的body
                Body invokeBody = SootCFG.getMethodBody(invokeExpr.getMethod().getDeclaringClass().getName(), invokeExpr.getMethod().getSignature());
                // 哪个参数发起的调用，方面重命名被调用方法里的变量，否则会重复
                Value caller = getCaller(jInvokeStmt.getInvokeExprBox().getValue());
                // 重命名被调用方法体里的变量
                Chain<Local> locals = invokeBody.getLocals();
                for (Local localVar : invokeBody.getLocals()) {
                    localVar.setName(caller + "_" + localVar.getName());
                }
                // 传给调用方法的参数
                List<Value> args = invokeExpr.getArgs();
                // 把调用的方法体里的自定的参数名字全都改成对应的参数
                for (int j = 0; j < args.size(); j++) {
                    invokeBody.getParameterLocal(j).setName(args.get(j).toString());
                }
                // 被调用函数的所有语句
                UnitPatchingChain invokeUnits = invokeBody.getUnits();
                // 被调用函数里把@this改为caller，将return语句换成caller赋值语句
                Iterator<Unit> iterator = invokeUnits.snapshotIterator();
                while (iterator.hasNext()) {
                    Unit invokeUnit = iterator.next();
                    // 有@this的赋值语句
                    if (invokeUnit instanceof JIdentityStmt) {
                        JIdentityStmt jiUnit = (JIdentityStmt) invokeUnit;
                        if (jiUnit.toString().contains("@this")) {
                            Value leftOp = jiUnit.getLeftOp();
                            ((Local) leftOp).setName(caller.toString());
                        }
                    }
                    // return语句
                    if (invokeUnit instanceof JReturnStmt || invokeUnit instanceof JReturnVoidStmt) {
                        // 注意调用结束后需要goto回调用的地方，也就是把这个return语句改成goto语句
                        JGotoStmt gotoStmt = new JGotoStmt(unitList.get(i + 1));
                        invokeUnits.insertBefore(gotoStmt, invokeUnit);
                        invokeUnits.remove(invokeUnit);
                    }
                }
                // 把这个调用方法的所有units插入原本的units，完成这个调用节点的扩展
                units.insertOnEdge(invokeUnits, unit, null);
            }

            // 在赋值语句中有调用
            if (unit instanceof JAssignStmt && ((JAssignStmt) unit).containsInvokeExpr()) {
                invokes.add(unit);
                JAssignStmt jAssignStmt = (JAssignStmt) unit;
                InvokeExpr invokeExpr = jAssignStmt.getInvokeExpr();
                Body invokeBody = SootCFG.getMethodBody(invokeExpr.getMethod().getDeclaringClass().getName(), invokeExpr.getMethod().getSignature());
                // 拿到赋值语句的左边，后面要为他赋值返回值
                Value ExternalLeftOp = jAssignStmt.getLeftOp();
                // 还是重命名
                Value caller = getCaller(jAssignStmt.getInvokeExprBox().getValue());
                Chain<Local> locals = invokeBody.getLocals();
                for (Local localVar : invokeBody.getLocals()) {
                    localVar.setName(caller + "_" + localVar.getName());
                }
                // 还是给参数赋值
                List<Value> args = invokeExpr.getArgs();
                for (int j = 0; j < args.size(); j++) {
                    invokeBody.getParameterLocal(j).setName(args.get(j).toString());
                }
                // 相同的处理
                UnitPatchingChain invokeUnits = invokeBody.getUnits();
                // 被调用函数里把@this改为caller，将return语句换成caller赋值语句
                Iterator<Unit> iterator = invokeUnits.snapshotIterator();
                while (iterator.hasNext()) {
                    Unit invokeUnit = iterator.next();
                    // 有@this的赋值语句
                    if (invokeUnit instanceof JIdentityStmt) {
                        JIdentityStmt jiUnit = (JIdentityStmt) invokeUnit;
                        if (jiUnit.toString().contains("@this")) {
                            Value leftOp = jiUnit.getLeftOp();
                            ((Local) leftOp).setName(caller.toString());
                        }
                    }
                    // return语句
                    if (invokeUnit instanceof JReturnStmt) {
                        // 这里有不同，要把返回值赋值给调用方法的这个赋值语句
                        JReturnStmt jreUnit = (JReturnStmt) invokeUnit;
                        Value returnVal = jreUnit.getOp();
                        JAssignStmt newAssignStmt = new JAssignStmt(ExternalLeftOp, returnVal);
                        invokeUnits.insertBefore(newAssignStmt, jreUnit);
                        // 注意调用结束后需要goto回调用的地方，也就是把这个return语句改成goto语句
                        JGotoStmt gotoStmt = new JGotoStmt(unitList.get(i + 1));
                        invokeUnits.insertBefore(gotoStmt, jreUnit);
                        invokeUnits.remove(jreUnit);
                    }
                }
                // 仍然是插入
                units.insertOnEdge(invokeUnits, unit, null);
            }
        }
    }

    // 获取调用语句中的调用者对象
    private static Value getCaller(Value jInvokeValue) {
        if (jInvokeValue instanceof JSpecialInvokeExpr) {
            JSpecialInvokeExpr jSpecialInvokeExpr = (JSpecialInvokeExpr) jInvokeValue;
            return jSpecialInvokeExpr.getBaseBox().getValue();
        } else if (jInvokeValue instanceof JVirtualInvokeExpr) {
            JVirtualInvokeExpr jVirtualInvokeExpr = (JVirtualInvokeExpr) jInvokeValue;
            return jVirtualInvokeExpr.getBaseBox().getValue();
        }
        return null; //TODO 目前只处理了两种invoke
    }

    // 默认根路径
    private static Body getMethodBody(String clsName, String metSignature) {
        G.reset();
        Options.v().set_prepend_classpath(true);
        Options.v().set_allow_phantom_refs(true);
        Options.v().set_soot_classpath(System.getProperty("user.dir") + File.separator + "target" + File.separator + "test-classes");
        SootClass sc = Scene.v().loadClassAndSupport(clsName);
        sc.setApplicationClass();
        Scene.v().loadNecessaryClasses();
        SootMethod sm = Scene.v().getMethod(metSignature);
        Body body = sm.retrieveActiveBody();
        return body;
    }
//    public static BriefUnitGraph getAdvancedCFG(String sourceDirectory,String clsName,String methodName) {
//        Body body = getMethodBody(sourceDirectory,clsName,methodName);
//        BriefUnitGraph adCfg = new BriefUnitGraph(body);
//        return adCfg;
//    }

    public static Body getMethodBody(String sourceDirectory,String clsName,String methodName){
        G.reset();
        // 设置在加载类时前面添加上set_soot_classpath里设置的类路径
        Options.v().set_prepend_classpath(true);
        // 允许加载虚拟引用。虚拟引用是 Soot 在分析中使用的一种引用，表示在源代码中存在但在字节码中缺失的类。
        Options.v().set_allow_phantom_refs(true);
        Options.v().set_soot_classpath(sourceDirectory);
        // 这一步将类加载到 Soot 中，但不一定会加载类的全部信息，直到后续需要。
        SootClass sc = Scene.v().loadClassAndSupport(clsName);
        // 标记加载的类 sc 为应用程序类，表明这是你关心的类。
        // 这意味着你关心这个类的详细信息，可能会在后续的静态分析、优化或者转换中使用它。
        // 标记一个类为应用程序类的目的在于在分析中集中关注用户代码，而忽略系统库等部分。
        // 设置了应用程序类后，Soot 在进行类加载和分析时，可能会更加关注这些类的详细信息。这对于一些静态分析工具非常有用，因为它们可以专注于用户自定义代码的分析，而不被系统库的复杂性所干扰。
        sc.setApplicationClass();
        // 加载必要的类。这一步通常会加载与目标类有关的其他类，以确保在分析和转换过程中所有必需的类都已加载。
        Scene.v().loadNecessaryClasses();
        // 从 Soot 场景中获取类 clsName 的 SootClass 对象，以便进一步操作。
        SootClass mainClass = Scene.v().getSootClass(clsName);
        // 从类 mainClass 中获取指定名称 methodName 的方法，应该是把它转换成了SootMethod。
        SootMethod sm = mainClass.getMethodByName(methodName);
        // 获取该方法的活动方法体（Body）。方法体是包含了方法的具体实现，包括局部变量、基本块等信息。
        Body body = sm.retrieveActiveBody();
        return body;
    }
}
